
import os
import sys

sys.path.append(os.path.split(sys.path[0])[0])
import shutil
from time import time
import SimpleITK as sitk
import scipy.ndimage as ndimage
from multiprocessing.dummy import Pool
import math
import numpy as np
'''
这里把图像维度调整到固定大小，通过调整spacing实现
'''

data_root = r"E:\AllData\zunYiBreast"  # 数据的根路径，可以设置为“nii_path”的上一级目录

volume_path = r"E:\AllData\zunYiBreast\Pretreatment96to3-32/volume"
seg_path = r"E:\AllData\zunYiBreast\Pretreatment96to3-32\segmentation"

pretreatment_volume_path = os.path.join(data_root, 'Pretreatment96to3-32-256/volume')
pretreatment_seg_path = os.path.join(data_root, 'Pretreatment96to3-32-256/segmentation')


def identicSpacing(file,processfile_path,save_path):
    sample_duration=32
    sample_size=256
    # 将CT和标签入读内存
    volume = sitk.ReadImage(os.path.join(processfile_path, file), sitk.sitkInt16)
    volume_array = sitk.GetArrayFromImage(volume)
    old_shape=volume_array.shape
    last_shape =np.array([sample_duration,sample_size,sample_size])#调整spacing之后的维度
    old_spacing=volume.GetSpacing()[2],volume.GetSpacing()[1],volume.GetSpacing()[0]
    old_spacing=np.array(old_spacing)
    dimension_adjustment=1/(old_shape/last_shape)
    volume_array_new = ndimage.zoom(volume_array, dimension_adjustment, order=3)
    '''
    后面要把spacing调好
    newspacing*newshape=oldshape*oldspacing
    '''
    new_volume = sitk.GetImageFromArray(volume_array_new)
    new_spacing=old_shape*old_spacing/volume_array_new.shape
    new_volume.SetDirection(volume.GetDirection())
    new_volume.SetOrigin(volume.GetOrigin())
    new_volume.SetSpacing((new_spacing[2], new_spacing[1], new_spacing[0]))
    print(file,new_volume.GetSpacing(),volume_array_new.shape)
    sitk.WriteImage(new_volume, os.path.join(save_path, file))
    # print("one successful!",file)




if __name__ == '__main__':
    shape=np.array([256,128,156])
    a=np.ones(shape)
    ad=np.array([256,64,64])/shape#new shape/old shape
    b=ndimage.zoom(a,ad , order = 3)
    print(b.shape)
    exit()

    # if os.path.exists(para.pretreatment_path):
    #     shutil.rmtree(para.pretreatment_path)
    if not os.path.exists(pretreatment_volume_path):
        os.makedirs(pretreatment_volume_path)
    if not os.path.exists(pretreatment_seg_path):
        os.makedirs(pretreatment_seg_path)
    start = time()
    file_list=os.listdir(volume_path)
    print("处理MRI")
    for file in file_list:
        identicSpacing(file,processfile_path = volume_path,save_path = pretreatment_volume_path)

    file_list = os.listdir(seg_path)
    print("处理标注")
    for file in file_list:
        identicSpacing(file, processfile_path = seg_path, save_path = pretreatment_seg_path)



